# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import numpy as np

import os
import sys

__dir__ = os.path.dirname(os.path.abspath(__file__))


sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '..')))

os.environ["FLAGS_allocator_strategy"] = 'auto_growth'

import cv2
import json
import paddle

from ppocr.data import create_operators, transform
from ppocr.modeling.architectures import build_model
from ppocr.postprocess import build_post_process
from ppocr.utils.save_load import init_model
from ppocr.utils.utility import get_image_file_list
import tools.program1 as program


def draw_det_res(dt_boxes, config, img, img_name):
    if len(dt_boxes) > 0:
        import cv2
        src_im = img
        for box in dt_boxes:
            box = box.astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
        save_det_path = os.path.dirname(config['Global'][
            'save_res_path']) + "/det_results/"
        if not os.path.exists(save_det_path):
            os.makedirs(save_det_path)
        save_path = os.path.join(save_det_path, os.path.basename(img_name))
        cv2.imwrite(save_path, src_im)
        #logger.info("The detected Image saved in {}".format(save_path))

def draw_det(dt_boxes, img):
    if len(dt_boxes) > 0:
        import cv2
        src_im = img
        for box in dt_boxes:
            box = box.astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
        return(src_im)
    else:
        return(img)


        #cv2.imwrite(save_path, src_im)
        #logger.info("The detected Image saved in {}".format(save_path))

def get_center(p1,p2,p3,p4):
    cx = int((p1[0]+p2[0]+p3[0]+p4[0])/4.0)
    cy = int((p1[1]+p2[1]+p3[1]+p4[1])/4.0)
    return [cx,cy]
def initmodel_det(det):

    global_config = det["config"]['Global']
    config = det['config']
    # build model
    model = build_model(config['Architecture'])

    init_model(det["config"], model, det["logger"])

    # build post process
    ##########################################
    #post_process_class = build_post_process(det['config']['PostProcess'])
    post_process_class = build_post_process(det['config']['PostProcess'], global_config)

    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name == 'KeepKeys':
            op[op_name]['keep_keys'] = ['image', 'shape']
        transforms.append(op)

    ops = create_operators(transforms, global_config)
    model.eval()
    return model, post_process_class, ops
def initmodel_cls(cls):

    global_config = cls["config"]['Global']
    config = cls['config']
    # build model
    model = build_model(config['Architecture'])

    init_model(cls["config"], model, cls["logger"])

    # build post process
    ##########################################
    #post_process_class = build_post_process(det['config']['PostProcess'])
    post_process_class = build_post_process(cls['config']['PostProcess'], global_config)

    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name == 'KeepKeys':
            op[op_name]['keep_keys'] = ['image']
        transforms.append(op)

    ops = create_operators(transforms, global_config)
    model.eval()
    return model, post_process_class, ops

def initmodel_rec(rec):
    config = rec["config"]
    global_config = config['Global']

    # build post process
    post_process_class = build_post_process(config['PostProcess'],
                                            global_config)

    # build model
    # build model
    if hasattr(post_process_class, 'character'):
        config['Architecture']["Head"]['out_channels'] = len(
            getattr(post_process_class, 'character'))

    model = build_model(config['Architecture'])
    init_model(config, model, rec["logger"])


    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name in ['RecResizeImg']:
            op[op_name]['infer_mode'] = True
        elif op_name == 'KeepKeys':
            if config['Architecture']['algorithm'] == "SRN":
                op[op_name]['keep_keys'] = [
                    'image', 'encoder_word_pos', 'gsrm_word_pos',
                    'gsrm_slf_attn_bias1', 'gsrm_slf_attn_bias2'
                ]
            else:
                op[op_name]['keep_keys'] = ['image']
        transforms.append(op)
    global_config['infer_mode'] = True
    ops = create_operators(transforms, global_config)
    model.eval()
    return model, post_process_class, ops




def rotate(ps, m):
    pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
    pts = np.hstack([pts, np.ones([len(pts), 1])]).T
    target_point = np.dot(m, pts)
    target_point = [[int(target_point[0][x]),int(target_point[1][x])] for x in range(len(target_point[0]))]
    return target_point

def rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
    h, w, c = img.shape
    M = cv2.getRotationMatrix2D((center_x,center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    out_points = rotate(points, M)
    return res_img, out_points

def rotate_img(img,angle,resize_rate=1.0):
    h, w, c = img.shape
    center_x = int(w/2.0)
    center_y = int(h/2.0)
    M = cv2.getRotationMatrix2D((center_x, center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    return res_img
def cls_(models,file):
    with open(file, 'rb') as f:
        img = f.read()
        data = {'image': img}
    batch = transform(data, models["cls"]["ops"])

    images = np.expand_dims(batch[0], axis=0)
    images = paddle.to_tensor(images)
    preds = models["cls"]["model"](images)
    post_result = models["cls"]["post_process_class"](preds)
    print(post_result)

def rec_(models, file1,file2):
    model = models['rec']['model']
    config = models['rec']["config"]
    global_config = config['Global']

    # build post process
    post_process_class = build_post_process(config['PostProcess'],
                                            global_config)

    # build model
    # build model
    if hasattr(post_process_class, 'character'):
        config['Architecture']["Head"]['out_channels'] = len(
            getattr(post_process_class, 'character'))


    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name in ['RecResizeImg']:
            op[op_name]['infer_mode'] = True
        elif op_name == 'KeepKeys':
            if config['Architecture']['algorithm'] == "SRN":
                op[op_name]['keep_keys'] = [
                    'image', 'encoder_word_pos', 'gsrm_word_pos',
                    'gsrm_slf_attn_bias1', 'gsrm_slf_attn_bias2'
                ]
            else:
                op[op_name]['keep_keys'] = ['image']
        transforms.append(op)
    global_config['infer_mode'] = True
    ops = create_operators(transforms, global_config)

    imglist = [file1, file2]
    results = []
    for file in imglist:
        with open(file, 'rb') as f:
            img = f.read()
            data = {'image': img}

        #config = models["rec"]["config"]
        batch = transform(data, ops)
        if config['Architecture']['algorithm'] == "SRN":
            encoder_word_pos_list = np.expand_dims(batch[1], axis=0)
            gsrm_word_pos_list = np.expand_dims(batch[2], axis=0)
            gsrm_slf_attn_bias1_list = np.expand_dims(batch[3], axis=0)
            gsrm_slf_attn_bias2_list = np.expand_dims(batch[4], axis=0)

            others = [
                paddle.to_tensor(encoder_word_pos_list),
                paddle.to_tensor(gsrm_word_pos_list),
                paddle.to_tensor(gsrm_slf_attn_bias1_list),
                paddle.to_tensor(gsrm_slf_attn_bias2_list)
            ]

        images = np.expand_dims(batch[0], axis=0)
        images = paddle.to_tensor(images)
        if config['Architecture']['algorithm'] == "SRN":
            preds = model(images, others)
        else:
            preds = model(images)
        post_result = post_process_class(preds)
        results.append(post_result)
    print(results)
    if len(results)>1:
        if results[0][0][1] > results[1][0][1]:
            return results[0][0][0]
        else:
            return results[1][0][0]


def test(imgpath, models):
    imglist = os.listdir(imgpath)
    for i in imglist:
        print(os.path.join(imgpath, i))
        with open(os.path.join(imgpath, i), 'rb') as f:
            img = f.read()
            data = {'image': img}
            f.close()
        batch = transform(data, models["det"]["ops"])
        images = np.expand_dims(batch[0], axis=0)
        shape_list = np.expand_dims(batch[1], axis=0)
        images = paddle.to_tensor(images)
        preds = models["det"]["model"](images)
        post_result = models["det"]["post_process_class"](preds, shape_list)
        boxes = post_result[0]['points']


        print(boxes)
        # write result
        dt_boxes_json = []
        for box in boxes:
            tmp_json = {}
            tmp_json['points'] = box.tolist()
            dt_boxes_json.append(tmp_json)
        # dt_boxes_json:[{'points': [[352, 189], [931, 614], [898, 691], [319, 266]]}]
        if len(dt_boxes_json)>0 and len(dt_boxes_json[0]['points']) > 0:
            biaoji = 0
            result = []
            if len(dt_boxes_json)>1:
                for dt in range(len(dt_boxes_json)):
                    sorts = dt_boxes_json[dt]['points']
                    sort_by = sorted(sorts, key=lambda point: point[1])
                    sorts = []
                    if sort_by[0][0] > sort_by[1][0]:
                        sorts.append(sort_by[1])
                        sorts.append(sort_by[0])
                    else:
                        sorts.append(sort_by[0])
                        sorts.append(sort_by[1])
                    if sort_by[2][0] > sort_by[3][0]:
                        sorts.append(sort_by[2])
                        sorts.append(sort_by[3])
                    else:
                        sorts.append(sort_by[3])
                        sorts.append(sort_by[2])
                        
                    result.append(distance(sorts[0],sorts[1])*distance(sorts[1],sorts[2]))

                for re in range(len(result)):
                    if result[re] > result[biaoji]:
                        biaoji = re
            print("result:",result)
            sorts = dt_boxes_json[biaoji]['points']
            sort_by = sorted(sorts, key=lambda point: point[1])
            sorts = []
            if sort_by[0][0] > sort_by[1][0]:
                sorts.append(sort_by[1])
                sorts.append(sort_by[0])
            else:
                sorts.append(sort_by[0])
                sorts.append(sort_by[1])
            if sort_by[2][0] > sort_by[3][0]:
                sorts.append(sort_by[2])
                sorts.append(sort_by[3])
            else:
                sorts.append(sort_by[3])
                sorts.append(sort_by[2])

            print(sorts)
            src_img = cv2.imread(os.path.join(imgpath, i))
            src_img = draw_det(boxes, src_img)

            #src_img = draw_det(boxes, src_img)
            sp = src_img.shape


            cv2.imshow('2', src_img)
            cv2.waitKey()
            cv2.destroyAllWindows()


        else:
            print("error")




            #src_img = cv2.imread(os.path.join(imgpath, i))
            #src_img = draw_det(boxes, src_img)



def distance(x,y):
    return  math.sqrt((x[0]-y[0])**2+(x[1]-y[1])**2)


def main(imgpath, config, device, logger, vdl_writer ):
    global_config = config['Global']

    # build model
    model = build_model(config['Architecture'])

    init_model(config, model, logger)

    # build post process
    post_process_class = build_post_process(config['PostProcess'])

    # create data ops
    transforms = []
    for op in config['Eval']['dataset']['transforms']:
        op_name = list(op)[0]
        if 'Label' in op_name:
            continue
        elif op_name == 'KeepKeys':
            op[op_name]['keep_keys'] = ['image', 'shape']
        transforms.append(op)

    ops = create_operators(transforms, global_config)
    model.eval()

    imglist = os.listdir(imgpath)
    for i in imglist:
        print(os.path.join(imgpath,i))
        with open(os.path.join(imgpath, i), 'rb') as f:
            img = f.read()
            data = {'image': img}
            f.close()
        batch = transform(data, ops)
        images = np.expand_dims(batch[0], axis=0)
        shape_list = np.expand_dims(batch[1], axis=0)
        images = paddle.to_tensor(images)
        preds = model(images)
        post_result = post_process_class(preds, shape_list)
        boxes = post_result[0]['points']
        print(boxes)
        # write result
        dt_boxes_json = []
        for box in boxes:
            tmp_json = {}
            tmp_json['points'] = box.tolist()
            dt_boxes_json.append(tmp_json)
        #dt_boxes_json:[{'points': [[352, 189], [931, 614], [898, 691], [319, 266]]}]
        if len(dt_boxes_json[0]['points']) > 0:
            sorts = dt_boxes_json[0]['points']
            sort_by = sorted(sorts, key=lambda point: point[1])
            sorts = []
            if sort_by[0][0] > sort_by[1][0]:
                sorts.append(sort_by[1])
                sorts.append(sort_by[0])
            else:
                sorts.append(sort_by[0])
                sorts.append(sort_by[1])
            if sort_by[2][0] >sort_by[3][0]:
                sorts.append(sort_by[2])
                sorts.append(sort_by[3])
            else:
                sorts.append(sort_by[3])
                sorts.append(sort_by[2])

            if distance(sorts[0],sorts[1]) > distance(sorts[1],sorts[2]):
                if sorts[1][1] > sorts[0][1]:
                    angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))
                    print(angel)
                else:
                    angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))
                    print("-", angel)
            else:
                if sorts[1][0] > sorts[2][0]:
                    angel = math.cos((sorts[1][0]-sorts[2][0])/distance(sorts[1],sorts[2]))
                    print(angel)
                else:
                    angel = math.cos((sorts[2][0]-sorts[1][0]) / distance(sorts[1], sorts[2]))
                    print("-",angel)







            print(sorts)

            src_img = cv2.imread(os.path.join(imgpath, i))
            src_img = draw_det(boxes, src_img)

            cv2.imshow('1', src_img)
            cv2.waitKey()
            cv2.destroyAllWindows()




if __name__ == '__main__':
    models = {"det":{}, "rec":{}, "cls":{}}
    #config, device, logger, vdl_writer = program.preprocess("../configs/det/det_r50_vd_db.yml")
    models["det"]["config"], models["det"]["device"], models["det"]["logger"], models["det"]["vdl_writer"] = program.preprocess("../configs/det/plate_det.yml")
    models["det"]["model"],models["det"]["post_process_class"],models["det"]["ops"] = initmodel_det(models['det'])

    #models["cls"]["config"], models["cls"]["device"], models["cls"]["logger"], models["cls"]["vdl_writer"] = program.preprocess("../configs/cls/cls_mv3.yml")
    #models["cls"]["model"], models["cls"]["post_process_class"],models["cls"]["ops"] = initmodel_cls(models['cls'])

    models["rec"]["config"], models["rec"]["device"], models["rec"]["logger"], models["rec"]["vdl_writer"] = program.preprocess("../configs/rec/rec_vin_train.yml")
    models["rec"]["model"], models["rec"]["post_process_class"], models["rec"]["ops"] = initmodel_rec(models['rec'])


    test('D:/BaiduNetdiskDownload/results_0604/04/bad/0111', models)
    #test('./g_sample_ori', models)
